#!/usr/bin/env python

import socket, time, subprocess, re, os, sys

class Speed_Tester(object):
    def __init__(self, ip=None, port=None):
        self.ip = []
        self.port = []
        
        if ip is not None and port is not None:
            if len(ip)>1 and len(port)>1:
                if len(ip)==len(port):
                    self.ip = ip[:]
                    self.port = port[:]
            else:
                self.ip.append(ip)
                self.port.append(port)
            
    
    def add_server(self, ip, port):
        if ip is None or port is None:
            return False
            
        self.ip.append(ip)
        self.port.append(port)
        return True
    
    def add_server_list(self, ip_list, port_list):
        if len(ip_list) != len(port_list):
            return False
        self.ip.extend(ip_list)
        self.port.extend(port_list)
        
        return True
    
    def clear_list(self):
        self.ip = []
        self.port = []
    
    def tcp_conn_test(self, timeout=10):
        '''
        Takes the time it takes to open a tcp connection
        Actually works pretty good
        '''
        if not self.ip or not self.port:
            return False
        
        timed = {}
        for i in range(0,len(self.ip)):
            try:
                the_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            except:
                print "Could not create socket"
                continue
            
            try:
                the_sock.settimeout(timeout)
                start = time.time()
                the_sock.connect((self.ip[i], self.port[i]))
                end = time.time()
                the_time = end-start
            except:
                the_sock.close()
                timed["%s:%d" % (self.ip[i], self.port[i])] = None
                continue
            timed["%s:%d" % (self.ip[i],self.port[i])] = the_time
            
        return timed
    
    def os_ping_test(self):
        '''
        Should give the best possible timing (takes an average), but takes longer time
        '''
        timed = {}
        for i in range(0,len(self.ip)):
            if sys.platform.startswith("win"):
                p = subprocess.Popen(["ping", "-n", "4", self.ip[i]], stdout = subprocess.PIPE)
                res = p.communicate()[0]
                
                reg = re.compile("Average = (\d*)ms", re.DOTALL)
                avg_time = re.search(reg, res)
                if avg_time is None:
                    timed["%s:%d" % (self.ip[i], self.port[i])] = None
                    continue
                timed["%s:%d" % (self.ip[i], self.port[i])] = avg_time.group(1)
            
            else:
                p = subprocess.Popen(["ping","-c", "4",self.ip[i]], stdout=subprocess.PIPE)
                res = p.communicate()[0]

                reg = re.compile("/(\d*\.\d{3})", re.DOTALL)
                avg_time = re.search(reg, res)
                if avg_time is None:
                    timed["%s:%d" % (self.ip[i], self.port[i])]=None
                    continue
                timed["%s:%d" % (self.ip[i], self.port[i])] = avg_time.group(1)
        
        return timed